import zarr
import numpy as np
import torch
from tqdm import tqdm
from franka_fk import init_ik_solver
assert zarr is not None, "zarr not imported"
assert np is not None, "numpy not imported"
assert torch is not None, "torch not imported"
def process_zarr_franka(zarr_path):
    zarr_data = zarr.open(zarr_path, mode='a')
    assert 'qpos' in zarr_data, f"qpos not found in {zarr_path}"
    assert 'state' in zarr_data['qpos'], f"state not found in qpos"
    qpos_state = zarr_data['qpos']['state']
    assert qpos_state.shape[1] == 16, f"Expected shape[1]=16, got {qpos_state.shape[1]}"
    ik_solver = init_ik_solver()
    timesteps = qpos_state.shape[0]
    eelocal_state = np.zeros((timesteps, 16))
    for timestep in tqdm(range(timesteps), desc="Processing timesteps"):
        left_joints = torch.tensor(qpos_state[timestep, 0:7], dtype=torch.float32).unsqueeze(0).cuda()
        right_joints = torch.tensor(qpos_state[timestep, 8:15], dtype=torch.float32).unsqueeze(0).cuda()
        left_kin_state = ik_solver.fk(left_joints)
        right_kin_state = ik_solver.fk(right_joints)
        left_pose = torch.cat([left_kin_state.ee_position, left_kin_state.ee_quaternion], dim=-1)
        right_pose = torch.cat([right_kin_state.ee_position, right_kin_state.ee_quaternion], dim=-1)
        eelocal_state[timestep, 0:7] = left_pose.cpu().numpy()
        eelocal_state[timestep, 7] = qpos_state[timestep, 7]
        eelocal_state[timestep, 8:15] = right_pose.cpu().numpy()
        eelocal_state[timestep, 15] = qpos_state[timestep, 15]
    if 'eelocal' not in zarr_data:
        zarr_data.create_group('eelocal')
    zarr_data['eelocal']['state'] = eelocal_state
    return zarr_data
if __name__ == "__main__":
    import sys
    assert len(sys.argv) == 2, f"Usage: {sys.argv[0]} <zarr_path>"
    zarr_path = sys.argv[1]
    result = process_zarr_franka(zarr_path)
    print(f"Processed {zarr_path}, added eelocal/state field") 